import os
import itertools
import argparse


def main():
    parser = argparse.ArgumentParser(description='MoE 超参数网格搜索驱动器')
    parser.add_argument('--train_cmd', type=str, default='python train_classifier.py', help='训练命令入口')
    parser.add_argument('--base_args', type=str, default='', help='固定基础参数（例如数据路径、类别数等）')
    parser.add_argument('--output_dir', type=str, default='runs_grid', help='输出目录根')
    args = parser.parse_args()

    # 超参数网格
    grids = {
        'lb_coeff': [0.01, 0.05, 0.1],
        'gating_strategy': ['softmax_topk', 'prob_norm'],
        'second_prob_threshold': [0.0, 0.05],
        'top_k': [2, 3],
        'visual_pretrained': [True, False],
        'use_visual_aux': [True, False],
        'visual_aux_weight': [0.3, 0.4, 0.5],
        'augment_train': ['basic', 'strong'],
        'rand_magnitude': [5, 7, 9],
    }

    keys = list(grids.keys())
    combos = list(itertools.product(*[grids[k] for k in keys]))
    print(f"总组合数: {len(combos)}")

    for i, combo in enumerate(combos, start=1):
        kv = dict(zip(keys, combo))
        run_dir = os.path.join(args.output_dir, f"run_{i:03d}")
        os.makedirs(run_dir, exist_ok=True)

        cmd_parts = [args.train_cmd]
        if args.base_args:
            cmd_parts.append(args.base_args)
        # 将布尔值转为flag
        for k, v in kv.items():
            if isinstance(v, bool):
                if v:
                    cmd_parts.append(f"--{k}")
            else:
                cmd_parts.append(f"--{k} {v}")
        cmd_parts.append(f"--output_dir {run_dir}")
        cmd = ' '.join(cmd_parts)
        print(f"[{i}/{len(combos)}] 运行: {cmd}")
        os.system(cmd)


if __name__ == '__main__':
    main()